#!/usr/bin/env python 

import os
import cmd
import sys
import random
import urllib2

from colors import * #@UnusedWildImport
from shelob import Shelob
from jpeg_analysis import JpegAnalysis
from statistics import BarGraph, Histogram, BoxPlot

def banner():
    quotes = ["It\'s not the figures lying, it\'s the liars figuring",
              "Statistics can be made to prove anything even the truth",
              "Facts are stubborn things, statistics are more pliable"]
    print "\t ,___________________________________________     "
    print "\t |\\'-._(   /                                 |   "
    print "\t | \\  .'-._\\  - Shelob Web Spider -     ,   ,|  "
    print "\t |-.\\'    .-;     Version: v0.2       .'\\`-' |  "
    print "\t |   \\  .' (                       _.'   \\   |  "
    print "\t |.--.\\'   _)  By: Moloch &     ;-;       \\._|  "
    print "\t |    ` _\\(_)/_         f47h3r   \\ `'-,_,-'\\ | "
    print "\t |______ /(O)\\  _________________/____)_`-._\\|  "
    quote = quotes[random.randint(0, len(quotes)-1)]
    print '\n    "'+quote+'"\n'
    
class ShelobUi(cmd.Cmd):

    intro = banner()
    prompt = underline+'shelob'+W+' > '
    doc_header = 'Spider Commands'
    undoc_header = 'Other'
    ruler = '='
    spider = Shelob()
    analysis = None
    
    def do_start(self, mode):
        ''' [*] Starts the web spider, optionally with a realtime monitor '''
        if mode == 'monitor':
            self.spider.startCollecting(True)
        else:
            self.spider.startCollecting()
        sys.stdout.write('\n')
    
    def do_set(self, userInput):
        ''' [*] Configure spider variables. Valid types are: file, timeout, pages, load, tmp '''
        try:
            type, value = userInput.split(' ')
            if type == 'file':
                self.spider.fileExtension = value
                print ' [*] Set file extension to:', value
            elif type == 'timeout':
                self.spider.timeout = int(value)
                print ' [*] Set timeout to:', value
            elif type == 'keyword':
                print ''
            elif type == 'inurl':
                print ''
            elif type == 'pages':
                self.spider.maxPages = int(value)
                print ' [*] Set maximum page count to:', value
            elif type == 'tmp':
                self.spider.setSaveDirectory(value)
            else:
                print C+' [-]'+W+' Not a valid type (%s), see help' % type
        except:
            print C+' [-]'+W+' Malformed input, see help'
            pass
    
    def do_seed(self, userInput):
        ''' [*] Add a search query to the root url list. Valid sites are: flickr, imageshack, photobucket, deviantArt'''
        try:
            site, seed = userInput.split(' ')
            seed = urllib2.quote(seed)
            if site == 'flickr':
                self.spider.addUrl('http://www.flickr.com/search/?q=%s&f=hp' % seed)
            elif site == 'imageshack':
                self.spider.addUrl('http://imageshack.us/photos/%s?' % seed)
            elif site == 'photobucket':
                self.spider.addUrl('http://photobucket.com/images/%s/' % seed)
            elif site == 'deviantArt':
                self.spider.addUrl('http://browse.deviantart.com/?qh=&section=&q=%s' % seed)
            elif site == '4chan':
                self.spider.addUrl('http://boards.4chan.org/%s' % seed)
                self.spider.addExclusion('thumbs')
            elif site == 'imgur':
                self.spider.addUrl('http://imgur.com/gallery/%s' % seed)
            elif site == 'wallbase':
                self.spider.addUrl('http://wallbase.cc/%s' % seed)
                self.spider.addExclusion('thumbs')
            else:
                print ' [*] Unknown (%s) site, see help' % site
        except:
            print C+' [-]'+W+' Malformed input, see help'
            pass
    
    def do_use(self, input):
        ''' [*] Select a search engine; Google, or Bing'''
        if input == "google":
            self.prompt = underline+'shelob'+W+R+' (Google)'+W+' > '
        elif input == "bing":
            self.prompt = underline+'shelob'+W+R+' (Bing)'+W+' > '
        else:
            print C+' [-]'+W+' Not a valid option, see help'
    
    def do_analyze(self, string):
        ''' [*] Parses all of the collected sample files for useful data '''
        if self.spider.saveDirectory == None:
            self.spider.setSaveDirectory()
        try:
            self.analysis = JpegAnalysis(self.spider.saveDirectory)
            self.analysis.analyzeFiles()
        except KeyError:
            print R+'\n [!] Warning: Possible corrupt file(s) in', self.spider.saveDirectory+W
            pass
    
    def do_add(self, url):
        ''' [*] Add any url to the spider's url list '''
        if url.startswith('http://'):
            self.spider.addUrl(url)
        else:
            print C+' [-]'+W+' Malformed url (%s), must start with "http://"' % url
    
    def do_exclude(self, userInput):
        ''' [*] Exclude image URLs that contain specified string(s) '''
        self.spider.addExclusion(userInput)
    
    def do_only(self, string):
        ''' [*] Only allow the spider to access URLs which contain specified strings '''
        self.spider.addExclusive(string)
    
    def do_bargraph(self, type):
        ''' [*] Generates a bar graph based on a metric.  Valid metrics are: size, date, '''
        if self.analysis == None:
            self.do_analyze('')
        data, title, xAxis, yAxis = self.analysis.getDictionary(type) #@UnusedVariable
        if data != None:
            graph = BarGraph(title, xAxis, yAxis)
            graph.buildGraph(data)
            graph.displayGraph()
        else:
            print ' [*] No data on %s, see help' % type
    
    def do_clear(self, type):
        ''' [*] Clear spider variables, valid types are: pages, files, excludes, all '''
        validTypes = ['pages', 'files', 'excludes', 'all']
        if type in validTypes:
            if type == 'pages' or type == 'all':
                self.spider.pages = []
                print ' [*] Cleared all spider pages'
            if type == 'files' or type == 'all':
                if self.spider.saveDirectory != None and os.path.exists(self.spider.saveDirectory):
                    os.remove(self.spider.saveDirectory)
                print ' [*] Removed all sample files'
            if type == 'excludes' or type == 'all':
                self.spider.exclude = []
                print ' [*] Cleared spider exclusion list'
        else:
            print C+' [-]'+W+' Malformed command, see help'
    
    def do_historgram(self, type):
        ''' [*] Generate a histogram for certain types '''
        if self.analysis == None:
            self.do_analyze('')
        data, title, xAxis, yAxis = self.analysis.getDictionary(type)
        if data != None:
            try:
                graph = Histogram(title, xAxis, yAxis)
                ls = []
                for key in data.keys():
                    for i in range(0, data[key]): #@UnusedVariable
                        ls.append(key)
                graph.buildGraph(ls)
                graph.displayGraph()
            except UnboundLocalError:
                print R+' [!] Error: Unable to generate a graph for this type'
                pass
            
    def do_boxplot(self, type):
        ''' [*] Generate a box and whiskers plot for certain types '''
        if self.analysis == None:
            self.do_analyze('')
        data, title, xAxis, yAxis = self.analysis.getDictionary(type)
        if data != None:
            try:
                graph = BoxPlot(title, xAxis, yAxis)
                ls = []
                for key in data.keys():
                    for i in range(0, data[key]): #@UnusedVariable
                        ls.append(key)
                graph.buildGraph(ls)
                graph.displayGraph()
            except TypeError:
                print R+' [!] Error: Unable to generate a graph for this type'+W
                pass
    
    def do_about(self, string):
        ''' [*] Display Shelob banner '''
        banner()
        
    def do_exit(self, string):
        ''' [*] Exit the Shelob interactive console '''
        print " [*] Good bye, have a wonderful day!"
        sys.exit()

if __name__ == '__main__':
    try:
        ShelobUi().cmdloop()
    except KeyboardInterrupt:
        print '\b\b\n [*] User exit, have a nice day!'
